-- Copyright (c) 2024 Huawei Technologies Co., Ltd.
-- openUBMC is licensed under Mulan PSL v2.
-- You can use this software according to the terms and conditions of the Mulan PSL v2.
-- You may obtain a copy of Mulan PSL v2 at: http://license.coscl.org.cn/MulanPSL2
-- THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
-- EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
-- MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
-- See the Mulan PSL v2 for more details.

local skynet = require 'skynet'
local log = require 'mc.logging'
local libmgmt_protocol = require 'libmgmt_protocol'
local class = require 'mc.class'

local gpu_object = class()

local properties = {
    'Capability0',
    'Capability1',
    'Capability2',
    'Capability3',
    'InfoRomVersion',
    'BoardPartNumber',
    'MemoryVendor',
    'MemoryPartNumber',
    'BuildDate',
    'UUID',
    'FirmwareVersion',
    'SN',
    'PartNumber',
    'ExternalPowerSufficient',
    'PowerBrakeSet',
    'PowerWatts',
    'ECCModeEnabled',
    'ECCModePendingEnabled',
    'ResetRequired'
}

function gpu_object:ctor(obj, total_power_sig)
    obj.Family = "N/A"
    obj.AssetType = "PCIe GPU Card"
    obj.InventorySerialNumber = obj.SerialNumber
    obj.InventoryFirmwareVersion = obj.CardFirmwareVersion
    obj.PCBVersion = "N/A"
    obj.InventoryManufacturer = obj.Manufacturer
    obj.AssetTag = "N/A"
    obj.InventoryPartNumber = obj.CardPartNumber
    obj.ManufactureDate = "N/A"
    obj.InventoryModel = obj.Model
    self.Capability0 = 0
    self.Capability1 = 0
    self.Capability2 = 0
    self.Capability3 = 0
    self.gpu = obj
    self.smbus_postbox_confg_obj = {}
    self.schedulers = {}
    self.metrics_sig = total_power_sig
end

function gpu_object:init_gpu_property()
    skynet.sleep(20 * 100)
    self.gpu.AssetName = string.format('GPU%s', self.gpu.Slot)
    self.gpu.InventorySlot = self.gpu.Slot
end

function gpu_object:init()
    skynet.fork(function()
        self:init_gpu_property()
    end)
end

function gpu_object:get_processor_info()
    return self.gpu
end

function gpu_object:clear_properties()
    for _, property in pairs(properties) do
        if property == 'PowerWatts' then
            self:set_prop(property, 0)
        elseif type(self.gpu[property]) == 'string' then
            self:set_prop(property, '')
        elseif type(self.gpu[property]) == 'number' then
            self:set_prop(property, 255)
        end
    end
end

function gpu_object:stop_nvlink_info()
    if self.scheduler_nvlink_info then
        self.scheduler_nvlink_info:deconstruct()
        self.scheduler_nvlink_info.data = nil
        self.scheduler_nvlink_info = nil
    end
end

function gpu_object:stop_action_task()
    if not self.action_task then
        return
    end
    self.action_task = false
    for _, s in ipairs(self.schedulers) do
        s:deconstruct()
        s.data = nil
    end
    self.schedulers = {}

    self:stop_nvlink_info()
    
end

function gpu_object:init_nvlink_info(nvlink_num)
    local MAX_NVLINK_NUM <const> = 256
    if nvlink_num == 0 or nvlink_num > MAX_NVLINK_NUM then
        return
    end

    self:stop_nvlink_info()

    local nvlink_info = self.smbus_postbox_confg_obj:NvLinkInfo()
    if nvlink_info and not nvlink_info.is_running then
        nvlink_info.params.request.ex_data_out = nvlink_num
        nvlink_info.on_data_change:on(function(data)
            local temp = {}
            for _, v in pairs(data) do
                -- T4不支持这4个参数，填充无效值
                table.insert(temp, {v, 32768, 32768, 32768, 32768})
            end
            self.gpu.NvLinkInfo = temp
        end)
        self.scheduler_nvlink_info = nvlink_info
        nvlink_info:start()
    end
end

function gpu_object:init_smbus_postbox()
    for _, property in pairs(properties) do
        local ok, s = pcall(self.smbus_postbox_confg_obj[property], self.smbus_postbox_confg_obj)
        if ok and s and not s.is_running then
            s.on_data_change:on(function(data)
                if self[property] then
                    self[property] = data
                else
                    self:set_prop(property, data)
                end
            end)
            table.insert(self.schedulers, s)
        end
    end

    -- RetiredPage
    local retired_page = self.smbus_postbox_confg_obj:RetiredPage()
    if retired_page and not retired_page.is_running then
        retired_page.on_data_change:on(function(data)
            self:set_prop('SingleBitErrorPageCount', data & 0xff)
            self:set_prop('DoubleBitErrorPageCount', (data & 0xff00) >> 8)
        end)
        table.insert(self.schedulers, retired_page)
    end

    -- NvLinkNum
    local nvlink_num = self.smbus_postbox_confg_obj:NvLinkNum()
    if nvlink_num and not nvlink_num.is_running then
        nvlink_num.on_data_change:on(function(data)
            -- 触发对NvLinkInfo的轮询
            self:init_nvlink_info(data)
        end)
        table.insert(self.schedulers, nvlink_num)
    end
    for _, s in ipairs(self.schedulers) do
        skynet.fork(s.start, s)
    end
end

function gpu_object:start_action_task()
    if self.action_task then
        return
    end
    self.action_task = true
    if next(self.smbus_postbox_confg_obj) then
        self:init_smbus_postbox()
        return
    end
    local ok, hardware_config = pcall(require, string.format('hardware_config.%s', self.gpu.Model:gsub(" ", "_")))
    if not ok or not hardware_config then
        log:notice('No specific hardware config file for this gpu(model: %s)', self.gpu.Model)
    else
        if hardware_config.smbus_postbox and self.gpu.RefChip then
            log:notice('hardware config file for this gpu(model: %s) loaded', self.gpu.Model)
            self.smbus_postbox_confg_obj =
                libmgmt_protocol.device_spec_parser(hardware_config.smbus_postbox(self.gpu.RefChip))
            self:init_smbus_postbox()
        end
    end
end

function gpu_object:set_prop(prop, val)
    if not self.gpu or not self.gpu[prop] then
        return
    end
    if self.gpu[prop] ~= val then
        if prop == 'PartNumber' and val == '' then
            log:debug('GPU %s is %s', prop, val)
            self:set_prop('PartNumber', 'Unknown')
        else
            self.gpu[prop] = val
        end
        if prop == 'SN' then
            self:set_prop('InventorySerialNumber', val)
        elseif prop == 'FirmwareVersion' then
            self:set_prop('InventoryFirmwareVersion', val)
        end

        if prop == 'PowerWatts' and self.metrics_sig then
            self.metrics_sig:emit()
        end
    end
end

local function convert_bdf(bdf)
    if not bdf then
        return
    end
    return string.match(bdf, '(.+):(.+):(.+).(.+)')
end

local BDF_PATH = 'bmc.kepler.sms.redfish.BDFNumber'
function gpu_object:bdf_match(gpu_resource)
    if not gpu_resource[BDF_PATH] or not gpu_resource[BDF_PATH].BDF then
        return
    end
    local _, dev_bus, dev_device, dev_func = convert_bdf(gpu_resource[BDF_PATH].BDF:value():value())
    if self.gpu.DevBus == tonumber(dev_bus, 16) and
        self.gpu.DevDevice == tonumber(dev_device, 16) and
        self.gpu.DevFunction == tonumber(dev_func, 16) then
        return true
    end
    return false
end

function gpu_object:transform_val(val)
    return tonumber(val:value():value()) or 0
end

function gpu_object:update_summary_info(gpu_info)
    for _ = 1, 3 do
        local ok, _ = pcall(self.update_summary, self, gpu_info)
        if ok then
            return
        end
        skynet.sleep(100)
    end
end

local UTILZATION_INTERFACE = 'bmc.kepler.sms.redfish.Utilization'
function gpu_object:update_summary(gpu_info)
    if not gpu_info.GPUResource and not gpu_info.Summary then
        return
    end
    if not self:bdf_match(gpu_info.GPUResource) then
        return
    end
    for _, o in pairs(gpu_info.Summary) do
        local utilization = o[UTILZATION_INTERFACE]
        self:set_prop('GPUUtilization', self:transform_val(utilization.GPU))
        self:set_prop('MemoryUtilization', self:transform_val(utilization.Memory))
    end
end

function gpu_object:update_summary_default()
    self:set_prop('GPUUtilization', 0)
    self:set_prop('MemoryUtilization', 0)
end

return gpu_object